# -*- encoding: utf-8 -*-
'''
@File    :   paderborn.py
@Time    :   2022/2/11 15:45
@Author  :   ZhangChaoYang
@Desc    :   德国帕德伯恩大学轴承数据集 http://groups.uni-paderborn.de/kat/BearingDataCenter/。由加速寿命测试引起的真实损坏：KA04,KA15,KB23,KB24,KI14,KI16
'''

import os, sys

import numpy as np

sys.path.insert(0, os.getcwd())
from scipy.io import loadmat
from util.work_flow import trans_collection, transform

fs = 64e3  # 采样频率64kHz
WC = ["N15_M07_F10", "N09_M07_F10", "N15_M01_F10", "N15_M07_F04"]
corpus_dir = 'E:\\zcy\\数据集\\Paderborn'
corpus_name = "paderborn"


def get_label(file):
    paths = os.path.split(os.path.dirname(file))
    if paths[-1].startswith("K0"):  # 最后一级目录以"K0"开头为正常样本
        return 1
    else:
        return 0


def get_wc(file):
    arr = file.split("_")
    return "_".join(arr[:3])


def gen_corpus(trans, dim):
    if trans not in trans_collection:
        raise Exception("unsupported transformation method: {}".format(trans))
    wc_ndarray = {wc: [None, None] for wc in WC}
    for bearing_dir in os.listdir(corpus_dir):
        bearing_dir = os.path.join(corpus_dir, bearing_dir)
        for basename in os.listdir(bearing_dir):
            if not basename.endswith(".mat"):
                continue
            wc = get_wc(basename)
            infile = os.path.join(bearing_dir, basename)
            label = get_label(infile)
            ndarray = wc_ndarray[wc][label]
            dic = loadmat(infile)
            for key in dic.keys():
                if key == basename[:-4]:
                    array = np.transpose(dic[key][0][0][2][0][6][2][0])  # 振动数据vibration_1
                    x = transform(trans, array, 4096, 1, fs)
                    if ndarray is None:
                        ndarray = x
                    else:
                        ndarray = np.vstack((ndarray, x))
                    wc_ndarray[wc][label] = ndarray

    for wc, ndarrays in wc_ndarray.items():
        for i, ndarray in enumerate(ndarrays):
            cls = 'normal' if i == 0 else 'anomaly'
            print(cls, wc, ndarray.shape)
            out_path = os.path.join("corpus", "multi_class", dim, corpus_name, trans, wc)
            if not os.path.exists(out_path):
                os.makedirs(out_path)
            np.save(os.path.join(out_path, cls), ndarray)


if __name__ == '__main__':
    gen_corpus('original', '1d')
    gen_corpus('fft', '1d')
    gen_corpus('stat', '1d')
    gen_corpus('sfft', '2d')
    # gen_corpus('cwt', '2d')
